{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Select nep.txt from UNEP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "select_symbol = ['W', 'Ta', 'V', 'Cr']\n",
    "n_max = np.array([6, 4])\n",
    "basis_size = np.array([12, 8])\n",
    "neuron = 80\n",
    "\n",
    "comment = []\n",
    "para = []\n",
    "with open('unep.txt', 'r') as fid:\n",
    "    lines = iter(fid.readlines())\n",
    "    for line in lines:\n",
    "        line = line.strip()\n",
    "        input_para = line.split()[0]\n",
    "        if input_para == 'nep4_zbl':\n",
    "            initial_symbol = line.split()[2:]\n",
    "            select_symbol = sorted(select_symbol, key=initial_symbol.index)\n",
    "            comment.append(f\"nep4_zbl {len(select_symbol)} \" + \" \".join(select_symbol))\n",
    "        else:\n",
    "            comment.append(line)\n",
    "        if input_para == 'ANN':\n",
    "            break\n",
    "    for line in lines:\n",
    "        para.append(float(line.split()[0]))\n",
    "\n",
    "para = np.array(para)\n",
    "N_des = n_max[0] + 1 + (n_max[1] + 1) * 5\n",
    "N_ann_one = (N_des + 2) * neuron\n",
    "N_c_radial = (n_max[0] + 1) * (basis_size[0] + 1) \n",
    "N_c_angular = (n_max[1] + 1) * (basis_size[1] + 1)\n",
    "\n",
    "with open(\"nep.txt\", \"w\") as fid:\n",
    "    for line in comment:\n",
    "        fid.write(f\"{line}\\n\")\n",
    "\n",
    "    offset = 0\n",
    "    for symbol in initial_symbol:\n",
    "        for _ in range(N_ann_one):\n",
    "            if symbol in select_symbol:\n",
    "                fid.write(f\"{para[offset]:15.7e}\\n\")\n",
    "            offset += 1\n",
    "\n",
    "    fid.write(f\"{para[offset]:15.7e}\\n\")\n",
    "    offset += 1\n",
    "\n",
    "    for m in range(N_c_radial + N_c_angular):\n",
    "        for n1 in initial_symbol:\n",
    "            for n2 in initial_symbol:\n",
    "                if n1 in select_symbol and n2 in select_symbol:\n",
    "                    fid.write(f\"{para[offset]:15.7e}\\n\")\n",
    "                offset += 1\n",
    "\n",
    "    for n in range(N_des):\n",
    "        fid.write(f\"{para[offset]:15.7e}\\n\")\n",
    "        offset += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit ZBL Parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy.optimize import curve_fit\n",
    "from wizard.io import read_xyz\n",
    "\n",
    "def quadratic(r, p1, p2, p3, p4, p5, p6):\n",
    "    Z1 = 74\n",
    "    Z2 = 74\n",
    "    a = 0.46848 / (Z1 ** 0.23 + Z2 ** 0.23) \n",
    "    A = 14.399645 * Z1 * Z2\n",
    "    x = r / a\n",
    "    return A / r * (p1 * np.exp(-p2 * x) + p3 * np.exp(-p4 * x) + p5 * np.exp(-p6 * x)) \n",
    "\n",
    "dimer = read_xyz('dimer.xyz')\n",
    "r = []\n",
    "y = []\n",
    "p0 = [0.5, 0.5, 0.5, 0.5, 0.5, 0.55]\n",
    "bounds = ([0.1, 0.1, 0.1, 0.1, 0.1, 0.1], [1, 1, 1, 1, 1, 1])\n",
    "for atoms in dimer:\n",
    "    r.append(np.linalg.norm(atoms[0].position - atoms[1].position))\n",
    "    y.append(atoms.info['energy'])\n",
    "r = np.array(r)\n",
    "y = np.array(y)\n",
    "\n",
    "popt, pcov = curve_fit(quadratic, r, y, p0=p0, bounds= bounds, maxfev=100000)\n",
    "print(popt)\n",
    "plt.scatter(r, y)\n",
    "plt.plot(r, quadratic(r, *popt), color='red')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
